import sqlite3
import os
from dotenv import load_dotenv
import json
from pydantic import BaseModel, Field
from langchain_core.tools import tool
from langchain_deepseek import ChatDeepSeek
from langgraph.prebuilt import create_react_agent
from pathlib import Path

# 加载环境变量
load_dotenv(override=True)
current_dir = Path(__file__).resolve().parent  # 获取当前脚本的绝对路径

# 读取提示词
prompts_file_path = current_dir / "agent_prompts.txt"  # 构建绝对路径
with open(prompts_file_path, "r", encoding="utf-8") as f:
    prompt = f.read()

# 创建SQL查询工具
description = """
当用户需要进行数据库查询工作时，请调用该函数。
该函数用sqlite运行一段SQL代码，完成数据查询相关工作，
本函数只负责运行SQL代码并进行数据查询，若要进行数据提取，则使用另一个extract_data函数。
"""

# 定义结构化参数模型
class SQLQuerySchema(BaseModel):
    sql_query: str = Field(description=description)

# 封装为 LangGraph 工具
@tool(args_schema=SQLQuerySchema)
def sql_inquire(sql_query: str) -> str:
    """
    当用户需要进行数据库查询工作时，请调用该函数。
    该函数用sqlite运行一段SQL代码，完成数据查询相关工作，
    本函数只负责运行SQL代码并进行数据查询。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对sqlite3中user_data数据库中各张表进行查询，并获得各表中的各类相关信息
    :return：sql_query在sqlite3中的运行结果。
    """
    
    # 加载环境变量
    load_dotenv(override=True)

    # 创建连接，数据库文件是database.db，如果不存在则会自动创建
    connection = sqlite3.connect('15000user_20250731.db')

    # 查询与检索数据
    cursor = connection.cursor()
    cursor.execute(sql_query)
    results = cursor.fetchall()
    # print("SQL 查询已成功执行，正在整理结果...")

    # 将结果以 JSON 字符串形式返回
    return json.dumps(results, ensure_ascii=False)

# 创建工具列表
tools = [sql_inquire]

# 创建模型
model = ChatDeepSeek(model="deepseek-chat")

# 创建Agent
graph = create_react_agent(model=model, tools=tools, prompt=prompt)
